#include <stdio.h>
#include <stdlib.h>
#include <pcap.h>
#include <netinet/ip.h>
#include <netinet/udp.h>
#include <arpa/inet.h>
#include <net/ethernet.h>

#define FILTER_SIZE 1024
#define PROMISC 1
#define TIMEOUT 1000

void packet_handler(u_char *user, const struct pcap_pkthdr *h, const u_char *bytes)
{
    struct ip *ip_header;
    struct udphdr *udp_header;
    char src_ip[INET_ADDRSTRLEN];
    char dst_ip[INET_ADDRSTRLEN];
    u_int ip_len;
    u_short src_port, dst_port;

    // 跳过以太网头（14字节）
    ip_header = (struct ip *)(bytes + sizeof(struct ether_header));
    ip_len = ip_header->ip_hl * 4; // IP头长度（32位字转换为字节）

    // 检查是否为UDP协议
    if (ip_header->ip_p != IPPROTO_UDP)
        return;

    // 获取UDP头
    udp_header = (struct udphdr *)(bytes + sizeof(struct ether_header) + ip_len);

    // 转换网络字节序到主机字节序
    src_port = ntohs(udp_header->uh_sport);
    dst_port = ntohs(udp_header->uh_dport);

    // 转换IP地址为字符串
    inet_ntop(AF_INET, &(ip_header->ip_src), src_ip, INET_ADDRSTRLEN);
    inet_ntop(AF_INET, &(ip_header->ip_dst), dst_ip, INET_ADDRSTRLEN);

    printf("[UDP Packet] %s:%d -> %s:%d | Size: %d bytes\n",
           src_ip, src_port, dst_ip, dst_port, h->len);
}

int main(int argc, char *argv[])
{
    bpf_u_int32 mask;
    bpf_u_int32 net;
    char *dev = NULL;
    pcap_t *handle;
    pcap_if_t *alldevs;
    pcap_if_t *d;
    char errbuf[PCAP_ERRBUF_SIZE];
    struct bpf_program fp;
    char filter_exp[FILTER_SIZE];
    if (argc < 3)
    {
        fprintf(stderr, "Usage: %s <target_ip> <target_port> [interface]\n", argv[0]);
        return 1;
    }

    // 获取网络设备
    if (argc > 3)
    {
        dev = argv[3];
    }
    else
    {
        // 使用 pcap_findalldevs 替代废弃的 pcap_lookupdev
        if (pcap_findalldevs(&alldevs, errbuf) == -1)
        {
            fprintf(stderr, "Couldn't find any devices: %s\n", errbuf);
            return 1;
        }

        // 使用第一个非环回设备
        for (d = alldevs; d != NULL; d = d->next)
        {
            if (!(d->flags & PCAP_IF_LOOPBACK))
            {
                dev = strdup(d->name);
                break;
            }
        }

        if (dev == NULL && alldevs != NULL)
        {
            // 如果没有非环回设备，使用第一个设备
            dev = strdup(alldevs->name);
        }

        if (dev == NULL)
        {
            fprintf(stderr, "No network devices found\n");
            pcap_freealldevs(alldevs);
            return 1;
        }

        pcap_freealldevs(alldevs);
    }

    printf("Using device: %s\n", dev);

    // 获取网络地址和掩码
    if (pcap_lookupnet(dev, &net, &mask, errbuf) == -1)
    {
        fprintf(stderr, "Couldn't get netmask for device %s: %s\n", dev, errbuf);
        net = 0;
        mask = 0;
    }

    const char *target_ip = argv[1];
    const char *target_port = argv[2];

    // 构造过滤表达式
    snprintf(filter_exp, FILTER_SIZE,
             "udp and (host %s) and (port %s)", target_ip, target_port);

    // 打开网络接口
    handle = pcap_open_live(dev, BUFSIZ, PROMISC, TIMEOUT, errbuf);
    if (handle == NULL)
    {
        fprintf(stderr, "Couldn't open device %s: %s\n", dev, errbuf);
        exit(1);
    }

    // 编译并设置过滤器
    if (pcap_compile(handle, &fp, filter_exp, 0, net) == -1)
    {
        fprintf(stderr, "Couldn't parse filter %s: %s\n", filter_exp, pcap_geterr(handle));
        exit(1);
    }

    if (pcap_setfilter(handle, &fp) == -1)
    {
        fprintf(stderr, "Couldn't install filter %s: %s\n", filter_exp, pcap_geterr(handle));
        exit(1);
    }

    printf("Capturing on %s: UDP %s:%s\n", dev, target_ip, target_port);

    // 开始捕获数据包（无限循环）
    pcap_loop(handle, -1, packet_handler, NULL);

    // 清理
    pcap_freecode(&fp);
    pcap_close(handle);
    if (argc <= 3)
        free(dev); // 释放我们分配的内存
    return 0;
}
